# -*- coding: utf-8 -*-
"""
Created on Fri Feb 21 15:59:27 2020

@author: lenovo03
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import talib

#本地读取数据
def ReadStockData():
    df = pd.read_csv('603688-2019.csv')
    df = df[['trade_date','close']]
    return df
df = ReadStockData()

p_list = []
date_list = []
for row in range(len(df)):
    close_price = df.loc[row,'close'] #float
    trade_date = df.loc[row,'trade_date'] #str
    p_list.append(close_price)
    date_list.append(trade_date)

dd = np.array(p_list) #numpy.ndarray

def detect_via_cusum_lg(ts, istart,threshold_times):
    """
    detect a time series using  cusum algorithm
    :param ts: the time series to be detected
    :param istart: the data from index 0 to index istart will be used as cold startup data to train 启动数据
    :param threshold_times: the times for setting threshold  设置阈值的时间
    :return:
    """
    
    S_h = 0  #high
    S_l = 0  #low
    S_list = np.zeros(istart)  #[0.,0.,0.]

    meanArray = talib.SMA(ts,timeperiod = istart)  #244,[nan,3.6755,3.6625,...]
    '''
    简介：简单移动平均线SMA，又称“算术移动平均线”，是指对特定期间的收盘价进行简单平均化
    用法：talib.SMA(数据, 周期)
    返回值：一位数组（numpy.ndarray）
    ''' 
    stdArray = talib.STDDEV(np.log(ts/meanArray),timeperiod = istart)  #标准差
    
    for i in range(istart, len(ts)):     #2----243
        tslog = np.log(ts[i] / meanArray[i - 1])

        S_h_ = max(0, S_h + tslog - stdArray[i-1])
        S_l_ = min(0, S_l + tslog + stdArray[i-1])

        if S_h_> threshold_times*stdArray[i-1]:
            S_list = np.append(S_list,1)
            print("{}:上涨信号".format(date_list[i]))
            S_h_ = 0
        elif abs(S_l_)> threshold_times*stdArray[i-1]:
            S_list = np.append(S_list, -1)
            print("{}:下跌信号".format(date_list[i]))
            S_l_ = 0
        else:
            S_list = np.append(S_list, 0)
        S_h = S_h_
        S_l = S_l_
    return S_list

listup,listdown = [],[]
s_list = detect_via_cusum_lg(dd,istart=2,threshold_times=10)  #244=242+2
#print(len(s_list))

for i in range(0,len(s_list)):
    if s_list[i] == 1:
        listup.append(i)
    elif s_list[i] == -1 :
        listdown.append(i)

total_up,total_down = 0,0
buy_point = []
sell_point = []
trade_datelist = []
trade_pricelist = []
trade_amountlist = []
trade_signallist = []
total_amount = 0

for i in range(0,len(s_list)):
    if s_list[i] == 1:
        total_up += 1 
        total_down = 0
    elif s_list[i] == -1:
        total_down += 1
        total_up = 0
    else:
        continue
    if total_down >= 2:
        #买入数量
        trade_amount = (total_down+2)*10
        trade_datelist.append(date_list[i])
        trade_pricelist.append(p_list[i])
        trade_amountlist.append(trade_amount)
        trade_signallist.append('buy')
        total_amount += trade_amount
        buy_point.append(i)
        print("{}买入{}个单位,持仓{}单位".format(date_list[i],trade_amount,total_amount))
    if total_up >= 2:
        if total_amount>0:
            #卖出数量
            trade_amount = (total_up-1)*10
            trade_datelist.append(date_list[i])
            trade_pricelist.append(p_list[i])
            trade_amountlist.append(trade_amount)
            trade_signallist.append('sell')
            sell_point.append(i)
            total_amount -= trade_amount
            print("{}卖出{}个单位,持仓{}单位".format(date_list[i],trade_amount,total_amount))
dl = {'date':trade_datelist,'price':trade_pricelist,'amount':trade_amountlist,'signal':trade_signallist}
trade_record = pd.DataFrame(dl)
print("交易记录：")
print(trade_record)
print("总交易次数{}".format(len(trade_record)))

#计算收益率
def calculateReturn():
    invest,income = 0,0
    for row in range(len(trade_record)):
        trade_price = trade_record.loc[row,'price'] #float
        trade_amount = trade_record.loc[row,'amount'] #str
        trade_signal = trade_record.loc[row,'signal']
        if trade_signal == 'buy':
            invest += trade_price*trade_amount
        else:
            income += trade_price*trade_amount
    return_rate = ((income+total_amount*p_list[-1])-invest)/invest
    return return_rate
        
print("投资回报率{}".format(calculateReturn()))        
print("上涨信号总计{}个".format(len(listup)))
print("下跌信号总计{}个".format(len(listdown)))

print(listup)
print(dd)
plt.figure(figsize=(10,5))    
plt.plot(dd, color='y', lw=2.)
plt.title(r'CUSUM趋势判断',fontproperties='SimHei',fontsize=20)
plt.xlabel('横轴：索引',fontproperties='SimHei',fontsize=15)
plt.ylabel('纵轴：价格',fontproperties='SimHei',fontsize=15)
plt.plot(dd, '^', markersize=2, color='r', label='UP signal', markevery=listup)
plt.plot(dd, 'v', markersize=2, color='g', label='DOWN signal', markevery=listdown)
for index in buy_point:
    plt.annotate(r'b',xy=(index,dd[index]),xytext=(index+0.5,dd[index]+0.2),arrowprops=dict(arrowstyle='->',connectionstyle='arc3,rad=.2'))
for index in sell_point:
    plt.annotate(r's',xy=(index,dd[index]),xytext=(index+0.5,dd[index]+0.2),arrowprops=dict(arrowstyle='->',connectionstyle='arc3,rad=.2'))
plt.legend()
plt.grid(True)
plt.show()